#!/usr/bin/env python3
"""
AIQuant系统启动脚本

这个脚本提供了一个统一的入口来启动AIQuant系统的各个组件。
支持不同的启动模式和配置选项。
"""

import os
import sys
import asyncio
import argparse
import logging
import signal
import subprocess
from pathlib import Path
from typing import Optional, List
import uvicorn
from rich.console import Console
from rich.panel import Panel
from rich.progress import Progress, SpinnerColumn, TextColumn
import time

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from aiquant.config.settings import get_settings
from aiquant.config.database import check_database_connection
from aiquant.config.redis_config import check_redis_connection

console = Console()
logger = logging.getLogger(__name__)

class SystemStarter:
    """系统启动器"""
    
    def __init__(self):
        self.settings = get_settings()
        self.processes: List[subprocess.Popen] = []
        self.shutdown_requested = False
        
    def setup_signal_handlers(self):
        """设置信号处理器"""
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
        
    def _signal_handler(self, signum, frame):
        """信号处理器"""
        console.print(f"\n[yellow]接收到信号 {signum}，正在关闭系统...[/yellow]")
        self.shutdown_requested = True
        self.shutdown()
        
    async def check_dependencies(self) -> bool:
        """检查系统依赖"""
        console.print("[blue]检查系统依赖...[/blue]")
        
        with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            console=console,
        ) as progress:
            # 检查数据库连接
            task1 = progress.add_task("检查数据库连接...", total=None)
            try:
                await check_database_connection()
                progress.update(task1, description="✅ 数据库连接正常")
            except Exception as e:
                progress.update(task1, description=f"❌ 数据库连接失败: {e}")
                return False
                
            # 检查Redis连接
            task2 = progress.add_task("检查Redis连接...", total=None)
            try:
                await check_redis_connection()
                progress.update(task2, description="✅ Redis连接正常")
            except Exception as e:
                progress.update(task2, description=f"❌ Redis连接失败: {e}")
                return False
                
            # 检查必要的目录
            task3 = progress.add_task("检查目录结构...", total=None)
            required_dirs = ['logs', 'cache', 'reports', 'data']
            for dir_name in required_dirs:
                dir_path = project_root / dir_name
                dir_path.mkdir(exist_ok=True)
            progress.update(task3, description="✅ 目录结构检查完成")
            
        return True
        
    def start_api_server(self, background: bool = False) -> Optional[subprocess.Popen]:
        """启动API服务器"""
        console.print("[green]启动API服务器...[/green]")
        
        cmd = [
            sys.executable, "-m", "uvicorn",
            "aiquant.main:app",
            "--host", self.settings.api_host,
            "--port", str(self.settings.api_port),
            "--workers", str(getattr(self.settings, 'api_workers', 1)),
        ]
        
        if self.settings.debug:
            cmd.append("--reload")
            
        if background:
            process = subprocess.Popen(
                cmd,
                cwd=project_root,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            self.processes.append(process)
            return process
        else:
            subprocess.run(cmd, cwd=project_root)
            return None
            
    def start_web_dashboard(self, background: bool = False) -> Optional[subprocess.Popen]:
        """启动Web仪表板"""
        console.print("[green]启动Web仪表板...[/green]")
        
        cmd = [
            sys.executable, "-m", "streamlit", "run",
            "aiquant/web/dashboard.py",
            "--server.port", str(getattr(self.settings, 'web_port', 8501)),
            "--server.address", getattr(self.settings, 'web_host', '0.0.0.0'),
            "--server.headless", "true",
        ]
        
        if background:
            process = subprocess.Popen(
                cmd,
                cwd=project_root,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            self.processes.append(process)
            return process
        else:
            subprocess.run(cmd, cwd=project_root)
            return None
            
    def start_celery_worker(self, background: bool = False) -> Optional[subprocess.Popen]:
        """启动Celery工作进程"""
        console.print("[green]启动Celery工作进程...[/green]")
        
        cmd = [
            sys.executable, "-m", "celery",
            "worker",
            "-A", "aiquant.tasks",
            "--loglevel=info",
            "--concurrency=4"
        ]
        
        if background:
            process = subprocess.Popen(
                cmd,
                cwd=project_root,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            self.processes.append(process)
            return process
        else:
            subprocess.run(cmd, cwd=project_root)
            return None
            
    def start_celery_beat(self, background: bool = False) -> Optional[subprocess.Popen]:
        """启动Celery调度器"""
        console.print("[green]启动Celery调度器...[/green]")
        
        cmd = [
            sys.executable, "-m", "celery",
            "beat",
            "-A", "aiquant.tasks",
            "--loglevel=info"
        ]
        
        if background:
            process = subprocess.Popen(
                cmd,
                cwd=project_root,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )
            self.processes.append(process)
            return process
        else:
            subprocess.run(cmd, cwd=project_root)
            return None
            
    async def start_all_services(self):
        """启动所有服务"""
        console.print(Panel.fit(
            "[bold green]AIQuant 量化投研系统[/bold green]\n"
            "[blue]正在启动所有服务...[/blue]",
            title="系统启动"
        ))
        
        # 检查依赖
        if not await self.check_dependencies():
            console.print("[red]依赖检查失败，无法启动系统[/red]")
            return False
            
        try:
            # 启动API服务器
            self.start_api_server(background=True)
            time.sleep(2)
            
            # 启动Web仪表板
            self.start_web_dashboard(background=True)
            time.sleep(2)
            
            # 启动Celery工作进程
            self.start_celery_worker(background=True)
            time.sleep(1)
            
            # 启动Celery调度器
            self.start_celery_beat(background=True)
            time.sleep(1)
            
            console.print(Panel.fit(
                f"[bold green]✅ 系统启动成功！[/bold green]\n\n"
                f"🌐 API服务器: http://{self.settings.api_host}:{self.settings.api_port}\n"
                f"📊 Web仪表板: http://{getattr(self.settings, 'web_host', '0.0.0.0')}:{getattr(self.settings, 'web_port', 8501)}\n"
                f"📚 API文档: http://{self.settings.api_host}:{self.settings.api_port}/docs\n\n"
                f"按 Ctrl+C 停止系统",
                title="启动完成"
            ))
            
            # 等待信号
            while not self.shutdown_requested:
                # 检查进程状态
                for process in self.processes[:]:
                    if process.poll() is not None:
                        console.print(f"[red]进程 {process.pid} 意外退出[/red]")
                        self.processes.remove(process)
                        
                await asyncio.sleep(1)
                
        except Exception as e:
            console.print(f"[red]启动失败: {e}[/red]")
            self.shutdown()
            return False
            
        return True
        
    def shutdown(self):
        """关闭所有服务"""
        console.print("[yellow]正在关闭服务...[/yellow]")
        
        for process in self.processes:
            try:
                process.terminate()
                process.wait(timeout=10)
            except subprocess.TimeoutExpired:
                process.kill()
            except Exception as e:
                console.print(f"[red]关闭进程时出错: {e}[/red]")
                
        self.processes.clear()
        console.print("[green]所有服务已关闭[/green]")
        
    def show_status(self):
        """显示系统状态"""
        console.print(Panel.fit(
            "[bold blue]AIQuant 系统状态[/bold blue]",
            title="状态检查"
        ))
        
        # 检查进程状态
        running_processes = []
        for process in self.processes:
            if process.poll() is None:
                running_processes.append(process.pid)
                
        if running_processes:
            console.print(f"[green]运行中的进程: {running_processes}[/green]")
        else:
            console.print("[yellow]没有运行中的进程[/yellow]")


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="AIQuant系统启动脚本")
    parser.add_argument(
        "command",
        choices=["start", "api", "web", "worker", "beat", "all", "status", "stop"],
        help="要执行的命令"
    )
    parser.add_argument(
        "--background", "-b",
        action="store_true",
        help="在后台运行"
    )
    parser.add_argument(
        "--debug", "-d",
        action="store_true",
        help="启用调试模式"
    )
    
    args = parser.parse_args()
    
    # 设置日志级别
    log_level = logging.DEBUG if args.debug else logging.INFO
    logging.basicConfig(
        level=log_level,
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
    )
    
    starter = SystemStarter()
    starter.setup_signal_handlers()
    
    try:
        if args.command == "start" or args.command == "all":
            await starter.start_all_services()
        elif args.command == "api":
            starter.start_api_server(background=args.background)
        elif args.command == "web":
            starter.start_web_dashboard(background=args.background)
        elif args.command == "worker":
            starter.start_celery_worker(background=args.background)
        elif args.command == "beat":
            starter.start_celery_beat(background=args.background)
        elif args.command == "status":
            starter.show_status()
        elif args.command == "stop":
            starter.shutdown()
            
    except KeyboardInterrupt:
        console.print("\n[yellow]用户中断，正在关闭...[/yellow]")
        starter.shutdown()
    except Exception as e:
        console.print(f"[red]启动失败: {e}[/red]")
        starter.shutdown()
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())